from langchain_core.messages import HumanMessage, AIMessage
from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI

from langgraph_coder.src.retriever.langchain_dataload import MyDataLoader

from typing_extensions import TypedDict

code_gen_prompt = ChatPromptTemplate.from_messages(
    [
        (
            "system",
            """您是一位具有LCEL（LangChain表达式语言）专业知识的编码助手。以下是一套完整的LCEL文档：
-------
{context}
-------
根据上述提供的文档回答用户问题。确保您提供的任何代码都可以在所有必需的导入和变量定义的情况下执行。
并且按照这个json格式输出
prefix: str - 问题描述及方法
imports: str - 代码块导入语句
code: str - 代码块不包括导入语句
以代码解决方案的描述来组织您的答案。然后列出导入。最后列出功能代码块。以下是用户问题：""",
        ),
        # placeholder占位符 表示下一行不是固定的消息
        ("placeholder", "{messages}"),
    ]
)

class Code(TypedDict):
    """关于LCEL问题的代码解决方案的样式"""
    prefix: str
    imports: str
    code: str

class GenerateChain:
    def __init__(self, dataloader: MyDataLoader, model_name: str = "qwen3-coder-480b-a35b-instruct"):
        self.llm = ChatOpenAI(base_url="https://dashscope.aliyuncs.com/compatible-mode/v1", model=model_name, temperature=0.1)
        self.dataloader = dataloader

    def get_solution(self, messages: list):
        context = self.dataloader.get_content()
        code_generate_chain = code_gen_prompt | self.llm.with_structured_output(Code)
        # 可以添加回退机制，用with_fallback
        solution = code_generate_chain.invoke(
            {
                "messages": messages,
                "context": context,
            }
        )
        print("大模型加载完成")
        print(type(solution), solution)
        return solution


if __name__ == '__main__':
    dataloader = MyDataLoader()
    chain1 = GenerateChain(dataloader)
    print(chain1.get_solution(HumanMessage(content="如何在LCEL中构建RAG链？")))



